import edu.princeton.cs.algs4.In;
import edu.princeton.cs.algs4.StdOut;

import java.util.HashSet;

public class BoggleSolver {

    private static final int[][] STEPS = {{-1, -1}, {0, -1}, {1, -1}, {-1, 0}, {1, 0}, {-1, 1}, {0, 1}, {1, 1}};
    private final Trie trie;

    // Initializes the data structure using the given array of strings as the dictionary.
    // (You can assume each word in the dictionary contains only the uppercase letters A through Z.)
    public BoggleSolver(String[] dictionary) {
        if (dictionary == null) {
            throw new IllegalArgumentException();
        }
        trie = new Trie(dictionary);
    }

    // Returns the set of all valid words in the given Boggle board, as an Iterable.
    public Iterable<String> getAllValidWords(BoggleBoard board) {
        HashSet<String> words = new HashSet<>();
        boolean[][] visited = new boolean[board.rows()][board.cols()];
        for (int row = 0; row < board.rows(); row++) {
            for (int col = 0; col < board.cols(); col++) {
                dfs(board, row, col, "", visited, words);
            }
        }
        return words;
    }

    private void dfs(BoggleBoard board, int row, int col, String word, boolean[][] visited, HashSet<String> words) {
        if (col < 0 || col >= board.cols() || row < 0 || row >= board.rows()) {
            return;
        }
        if (visited[row][col]) return;
        visited[row][col] = true;
        char c = board.getLetter(row, col);
        word += c;
        if (c == 'Q') {
            word += 'U';
        }
        int state = trie.contains(word);
        if (word.length() > 2 && state > 0) {
            words.add(word);
        }
        if (state >= 0) {
            for (int[] st : STEPS) {
                dfs(board, row + st[0], col + st[1], word, visited, words);
            }
        }
        visited[row][col] = false;
    }

    // Returns the score of the given word if it is in the dictionary, zero otherwise.
    // (You can assume the word contains only the uppercase letters A through Z.)
    public int scoreOf(String word) {
        if (word == null) {
            throw new IllegalArgumentException();
        }
        if (trie.contains(word) > 0) {
            int length = word.length();

            switch (length) {
                case 0:
                case 1:
                case 2:
                    return 0;
                case 3:
                case 4:
                    return 1;
                case 5:
                    return 2;
                case 6:
                    return 3;
                case 7:
                    return 5;
                case 8:
                default:
                    return 11;
            }
        }
        return 0;
    }

    private static class Trie {
        private static final int R = 26;
        private final Node root = new Node();

        private static class Node {
            private Node[] next;
            private boolean isString;

            private Node() {
                this.next = new Node[R];
            }
        }

        Trie(String[] dict) {
            for (String s : dict) {
                add(root, s, 0);
            }
        }

        public int contains(String key) {
            if (key == null) {
                throw new IllegalArgumentException("argument to contains() is null");
            } else {
                Node x = this.get(this.root, key, 0);
                return x == null ? -1 : (x.isString ? 1 : 0);
            }
        }

        private Node get(Node x, String key, int d) {
            if (x == null) {
                return null;
            } else if (d == key.length()) {
                return x;
            } else {
                char c = key.charAt(d);
                return this.get(x.next[c - 'A'], key, d + 1);
            }
        }

        private Node add(Node x, String key, int d) {
            if (x == null) {
                x = new Node();
            }

            if (d == key.length()) {
                x.isString = true;
            } else {
                char c = key.charAt(d);
                x.next[c - 'A'] = this.add(x.next[c - 'A'], key, d + 1);
            }
            return x;
        }
    }

    public static void main(String[] args) {
        String path = args[0];
        In in = new In(path + "dictionary-yawl.txt");
        String[] dic = in.readAllLines();
        BoggleSolver boggleSolver = new BoggleSolver(dic);
        BoggleBoard boggleBoard = new BoggleBoard(path + "board-dichlorodiphenyltrichloroethanes.txt");

        int score = 0;
        int count = 0;
        for (String word : boggleSolver.getAllValidWords(boggleBoard)) {
            score += boggleSolver.scoreOf(word);
            StdOut.println(word);
            count++;
        }
        StdOut.println("Count is: " + count);
        StdOut.println("Score is: " + score);
        StdOut.println(boggleBoard);
    }
}
